import numpy as np

class LinearRegression:
    def __init__(self, lr=0.01, epochs=1000, batch_size=32):
        """
        初始化线性回归模型的超参数
        参数:
        lr (float): 学习率
        epochs (int): 训练轮数
        batch_size (int): 批次大小
        """
        self.lr = lr
        self.epochs = epochs
        self.batch_size = batch_size
        self.w = None
        self.b = None

    def fit(self, X, y):
        """
        训练线性回归模型
        参数:
        X (ndarray): 训练数据集，形状为 (n_samples, n_features)
        y (ndarray): 训练标签，形状为 (n_samples,)
        """
        m, n = X.shape  # 获取数据的维度
        self.w = np.zeros(n)  # 初始化权重
        self.b = 0  # 初始化偏置

        # 梯度下降过程
        for epoch in range(self.epochs):
            permutation = np.random.permutation(m)
            X = X[permutation]
            y = y[permutation]

            for i in range(0, m, self.batch_size):
                X_batch = X[i:i + self.batch_size]
                y_batch = y[i:i + self.batch_size]

                predictions = np.dot(X_batch, self.w) + self.b

                # 计算损失 (均方误差损失)
                loss = np.mean((predictions - y_batch) ** 2)

                # 计算梯度
                dw = (2 / self.batch_size) * np.dot(X_batch.T, (predictions - y_batch))
                db = (2 / self.batch_size) * np.sum(predictions - y_batch)

                # 更新权重和偏置
                self.w -= self.lr * dw
                self.b -= self.lr * db

            if epoch % 100 == 0:
                print(f"Epoch {epoch}, Loss: {loss}")

    def predict(self, X):
        """
        使用训练好的模型进行预测
        参数:
        X (ndarray): 输入数据，形状为 (n_samples, n_features)

        返回:
        ndarray: 预测结果（连续值）
        """
        return np.dot(X, self.w) + self.b

    def score(self, X, y):
        """
        计算模型的R^2得分
        参数:
        X (ndarray): 测试数据集，形状为 (n_samples, n_features)
        y (ndarray): 测试标签，形状为 (n_samples,)

        返回:
        float: R^2得分
        """
        y_pred = self.predict(X)
        ss_res = np.sum((y - y_pred) ** 2)
        ss_tot = np.sum((y - np.mean(y)) ** 2)
        r2_score = 1 - (ss_res / ss_tot)
        return r2_score